Add Esm #2244

pass-lin · 2025-05-03T07:40:33Z

from #2177
Achieved a smaller error with hf.

import os
os.environ["KERAS_BACKEND"] = "torch"
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

from keras import ops
from transformers.models.esm.modeling_esm import EsmAttention as hf_EsmSelfAttention
from transformers import EsmConfig
from esm2.esm2_layers import EsmSelfAttention
import numpy as np
import keras
from transformers.models.esm.modeling_esm import EsmModel
weights_path = "facebook/esm2_t6_8M_UR50D"
hf_model = EsmModel.from_pretrained(weights_path)
hf_model.cuda().eval()
hf_model.embeddings.token_dropout = False


from keras_hub.src.models.esm.esm_backbone import (
    ESMBackbone,
)


keras_model =  ESMBackbone.from_preset('hf://'+weights_path)
keras_model.summary()


x = ops.array([[1,2,3,4,5]])+1
hf_out = hf_model(x,ops.ones_like(x))[0]
keras_out = keras_model({'token_ids': x})

print(ops.all(ops.isclose(hf_out, keras_out,atol=1e-4)))

pass-lin · 2025-05-03T07:56:28Z

ruff.....................................................................Passed
ruff-format..............................................................Passed
Error: Process completed with exit code 1.

Please help me figure out how to solve this problem.

mattdangerw · 2025-05-06T18:35:36Z

Probably an issue with generating the API symbols. Looks like you need to sync with the latest changes on master, then you could try running ./shell/api_gen.sh

sachinprasadhs · 2025-05-09T17:15:23Z

ruff.....................................................................Passed
ruff-format..............................................................Passed
Error: Process completed with exit code 1.

Please help me figure out how to solve this problem.

You can rebase it to latest master code
and then run - pre-commit run --all-files
pip install -u namex

pass-lin · 2025-05-10T13:13:18Z

keras_hub/src/layers/modeling/reversible_embedding_test.py::ReversibleEmbeddingTest::test_quantize_dtype_argument_tie_weights - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/layers/modeling/reversible_embedding_test.py::ReversibleEmbeddingTest::test_quantize_dtype_argument_untie_weights - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/layers/modeling/reversible_embedding_test.py::ReversibleEmbeddingTest::test_quantize_int8_tie_weights - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/layers/modeling/reversible_embedding_test.py::ReversibleEmbeddingTest::test_quantize_int8_untie_weights - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/albert/albert_backbone_test.py::AlbertBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/bart/bart_backbone_test.py::BartBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/bert/bert_backbone_test.py::BertBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/bloom/bloom_backbone_test.py::BloomBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/clip/clip_backbone_test.py::CLIPBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/deberta_v3/deberta_v3_backbone_test.py::DebertaV3BackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/distil_bert/distil_bert_backbone_test.py::DistilBertBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/electra/electra_backbone_test.py::ElectraBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/f_net/f_net_backbone_test.py::FNetBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/falcon/falcon_backbone_test.py::FalconBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/gemma/gemma_backbone_test.py::GemmaBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/gemma/gemma_backbone_test.py::Gemma2BackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/gpt2/gpt2_backbone_test.py::GPT2BackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/gpt_neo_x/gpt_neo_x_backbone_test.py::GPTNeoXBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/llama/llama_backbone_test.py::LlamaTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/mistral/mistral_backbone_test.py::MistralBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/opt/opt_backbone_test.py::OPTBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/pali_gemma/pali_gemma_backbone_test.py::PaliGemmaBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/pali_gemma/pali_gemma_backbone_test.py::PaliGemma2BackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/phi3/phi3_backbone_test.py::Phi3Test::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/phi3/phi3_backbone_test.py::Phi3Test::test_backbone_basics_with_su_rotary - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/roberta/roberta_backbone_test.py::RobertaBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/siglip/siglip_backbone_test.py::SigLIPBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/siglip/siglip_backbone_test.py::SigLIP2BackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/t5/t5_backbone_test.py::T5BackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/whisper/whisper_backbone_test.py::WhisperBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/xlm_roberta/xlm_roberta_backbone_test.py

@mattdangerw @sachinprasadhs
Is it a problem with the test environment? Why are there so many errors that don't belong to me?

sachinprasadhs · 2025-05-12T17:44:15Z

It's not related to your code, looks like some issue with the JAX backend, we will look into it.

sachinprasadhs

Thanks fro the PR, I have added my comments, also add checkpoints conversion under: keras-hub/tools/checkpoint_conversion

keras_hub/src/models/esm/esm_backbone.py

sachinprasadhs · 2025-05-15T22:02:55Z

keras_hub/src/models/esm/esm_backbone.py

+        intermediate_dim: int. The output dimension of the first Dense layer in
+            a two-layer feedforward network for each transformer.
+        dropout: float. Dropout probability for the Transformer encoder.
+        layer_norm_eps:bool.Should we use ln after embedding?


Didn't get the point here, are you asking our input or it's the arg detail, if it is the arg details, it needs to be repharsed, avoid question marks and the argument name is emb_layer_norm_before

layer_norm_eps discription needs to be updated.

keras_hub/src/models/esm/esm_backbone.py

keras_hub/src/models/esm/esm_classifier_test.py

keras_hub/src/models/esm/esm_masked_plm_test.py

keras_hub/src/models/esm/esm_masked_plm.py

keras_hub/src/models/esm/esm_classifier.py

keras_hub/src/utils/transformers/convert_esm.py

pass-lin · 2025-05-17T18:13:05Z

@sachinprasadhs @mattdangerw
Can anybody review my code?

pass-lin · 2025-06-02T18:11:23Z

@mattdangerw @sachinprasadhs
Please check my code, thank you.

sachinprasadhs

Added few more comments and few of the previous review comments still needs to be addressed

sachinprasadhs · 2025-06-02T18:26:41Z

keras_hub/src/models/esm/esm_backbone.py

+        layer_norm_eps:bool.If true, then layer norm will be used before
+                        entering the transformer block.
+                        Since it's pre-norm, the default is false.


This is not bool as per the usage I can see, did you mean someother argument?

inconsitent indentation in args, follow 4 space indentation.

sachinprasadhs · 2025-06-02T18:37:34Z

keras_hub/src/models/esm/esm_backbone.py

+    Disclaimer: Pre-trained models are provided on an "as is" basis, without
+    warranties or conditions of any kind.
+
+    Args:


Still activation and max_wavelength description is missing!

sachinprasadhs · 2025-06-02T18:41:15Z

keras_hub/src/models/esm/esm_backbone.py

+        intermediate_dim: int. The output dimension of the first Dense layer in
+            a two-layer feedforward network for each transformer.
+        dropout: float. Dropout probability for the Transformer encoder.
+                    Defaults to 0.1


only 4 space indentaion

sachinprasadhs · 2025-06-02T18:42:21Z

keras_hub/src/models/esm/esm_backbone.py

+    Disclaimer: Pre-trained models are provided on an "as is" basis, without
+    warranties or conditions of any kind.
+
+    Args:


add arg description for pad_token_id as well

sachinprasadhs · 2025-06-02T18:46:05Z

keras_hub/src/models/esm/esm_backbone.py

+            embeddings.
+        position_embedding_type:esm1 use abs position embeding,esm2 use rope.
+            so this parameter is only except for absolute and rotary.
+        dtype: None or str or .keras.mixed_precision.DTypePolicy. The dtype to


fix typo .keras.mixed_precision.DTypePolicy --> keras.mixed_precision.DTypePolicy

sachinprasadhs · 2025-06-02T18:46:33Z

keras_hub/src/models/esm/esm_backbone.py

+        position_embedding_type:esm1 use abs position embeding,esm2 use rope.
+            so this parameter is only except for absolute and rotary.


This still needs to be changed to:

position_embedding_type: str. The position embedding type to use. One of "absolute" and "rotary". Use "absolute" for ESM1. Use "rotary" for ESM2. Defaults to "rotary".

sachinprasadhs · 2025-06-02T18:48:10Z

keras_hub/src/models/esm/esm_backbone_test.py

+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+            expected_output_shape=(2, 5, 2),
+        )


Still missing save model test

sachinprasadhs · 2025-06-02T18:50:08Z

keras_hub/src/models/esm/esm_classifier_preprocessor.py

+
+
+@keras_hub_export("keras_hub.models.ESMProteinClassifierPreprocessor")
+class ESMProteinClassifierPreprocessor(BertTextClassifierPreprocessor):


Pending change here which should be subclassed from TextClassifierPreprocessor instead of BertTextClassifierPreprocessor

sachinprasadhs · 2025-06-02T18:55:23Z

keras_hub/src/models/esm/esm_backbone.py

+        max_sequence_length=1024,
+        max_wavelength=10000,
+        layer_norm_eps=1e-12,
+        emb_layer_norm_before=False,


pending change, instead emb_layer_norm_before --> use_pre_layer_norm

sachinprasadhs · 2025-06-02T18:56:17Z

keras_hub/src/models/esm/esm_classifier.py

+
+
+@keras_hub_export("keras_hub.models.ESMProteinClassifier")
+class ESMProteinClassifier(RobertaTextClassifier):


pending change.
You can subclass TextClassifier and make the same changes as RobertaTextClassifier instead of subclassing from another model.

sachinprasadhs · 2025-06-02T19:13:43Z

Once you address all the comments, add end to end working colab along with the checkpoints conversion under: keras-hub/tools/checkpoint_conversion

pass-lin · 2025-06-03T10:53:20Z

Once you address all the comments, add end to end working colab along with the checkpoints conversion under: keras-hub/tools/checkpoint_conversion

Ok, please check the new code.

pass-lin added 2 commits May 3, 2025 01:28

add esm

f9ff098

add esm2

cc4123b

pass-lin added 4 commits May 3, 2025 17:35

fix

d3f598d

fix

737a147

format

140207b

fix test

cc9a11c

divyashreepathihalli requested a review from sachinprasadhs May 5, 2025 17:09

format

f8da784

pass-lin and others added 2 commits May 10, 2025 19:26

renew

72e9829

Merge branch 'keras-team:master' into esm

16bb9f2

pass-lin force-pushed the esm branch from a66ee78 to 19f4b1f Compare May 10, 2025 11:33

format

5cbf577

pass-lin force-pushed the esm branch from 19f4b1f to 5cbf577 Compare May 10, 2025 12:05

format

6e9f817

pass-lin mentioned this pull request May 10, 2025

_int8_build() bug from keras-nightly keras-team/keras#21272

Closed

sachinprasadhs reviewed May 16, 2025

View reviewed changes

pass-lin added 3 commits May 17, 2025 12:18

update

2815e9c

update

79e738c

update

20d5051

sachinprasadhs reviewed Jun 2, 2025

View reviewed changes

pass-lin added 2 commits June 3, 2025 17:59

update

fb18c98

add new tool

7609ab4

		position_embedding_type:esm1 use abs position embeding,esm2 use rope.
		so this parameter is only except for absolute and rotary.



		@keras_hub_export("keras_hub.models.ESMProteinClassifierPreprocessor")
		class ESMProteinClassifierPreprocessor(BertTextClassifierPreprocessor):



		@keras_hub_export("keras_hub.models.ESMProteinClassifier")
		class ESMProteinClassifier(RobertaTextClassifier):

Add Esm #2244

Are you sure you want to change the base?

Add Esm #2244

Uh oh!

Conversation

pass-lin commented May 3, 2025

Uh oh!

pass-lin commented May 3, 2025

Uh oh!

mattdangerw commented May 6, 2025

Uh oh!

sachinprasadhs commented May 9, 2025

Uh oh!

pass-lin commented May 10, 2025

Uh oh!

sachinprasadhs commented May 12, 2025

Uh oh!

sachinprasadhs left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

sachinprasadhs May 15, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

pass-lin commented May 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pass-lin commented Jun 2, 2025

Uh oh!

sachinprasadhs left a comment

Choose a reason for hiding this comment

Uh oh!

sachinprasadhs Jun 2, 2025

Choose a reason for hiding this comment

Uh oh!

sachinprasadhs Jun 2, 2025

Choose a reason for hiding this comment

Uh oh!

sachinprasadhs Jun 2, 2025

Choose a reason for hiding this comment

Uh oh!

sachinprasadhs Jun 2, 2025

Choose a reason for hiding this comment

Uh oh!

sachinprasadhs Jun 2, 2025

Choose a reason for hiding this comment

Uh oh!

sachinprasadhs Jun 2, 2025

Choose a reason for hiding this comment

Uh oh!

sachinprasadhs Jun 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sachinprasadhs Jun 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sachinprasadhs Jun 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sachinprasadhs Jun 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sachinprasadhs Jun 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sachinprasadhs commented Jun 2, 2025

Uh oh!

pass-lin commented May 17, 2025 •

edited

Loading

sachinprasadhs Jun 2, 2025 •

edited

Loading

sachinprasadhs Jun 2, 2025 •

edited

Loading

sachinprasadhs Jun 2, 2025 •

edited

Loading

sachinprasadhs Jun 2, 2025 •

edited

Loading

sachinprasadhs Jun 2, 2025 •

edited

Loading